{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "629c902e-6cd0-4475-b6ce-5d6e37d7e2f3",
   "metadata": {},
   "source": [
    "# T5 LLM Fine-Tuning with DeepSpeed and Kubeflow Trainer\n",
    "\n",
    "\n",
    "This Notebook will fine-tune Text-to-Text Transfer Transformer (T5) with CNN Dailymail dataset for text summarization using Kubeflow TrainJob and DeepSpeed.\n",
    "\n",
    "Pretrained T5 model: https://huggingface.co/google-t5/t5-base\n",
    "\n",
    "CNN Dailymail dataset: https://huggingface.co/datasets/abisee/cnn_dailymail\n",
    "\n",
    "This Notebook will use **4 x A100 NVIDIA GPUs**, to fine-tune T5 model on 2 nodes (every node has 4 GPUs)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c461b2984e77d99",
   "metadata": {},
   "source": [
    "## Install the Kubeflow SDK\n",
    "\n",
    "You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4900404c5d532bdf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:07:19.635245Z",
     "iopub.status.busy": "2025-10-18T00:07:19.634526Z",
     "iopub.status.idle": "2025-10-18T00:07:19.643133Z",
     "shell.execute_reply": "2025-10-18T00:07:19.642289Z",
     "shell.execute_reply.started": "2025-10-18T00:07:19.635202Z"
    }
   },
   "outputs": [],
   "source": [
    "# !pip install -U kubeflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47534ee4955f3ff6",
   "metadata": {},
   "source": [
    "## Create Script to Fine-Tune T5 with DeepSpeed\n",
    "\n",
    "We need to wrap our fine-tuning script into a function to create Kubeflow TrainJob."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35f06c45b614ecd0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:07:22.323778Z",
     "iopub.status.busy": "2025-10-18T00:07:22.323527Z",
     "iopub.status.idle": "2025-10-18T00:07:22.347076Z",
     "shell.execute_reply": "2025-10-18T00:07:22.346538Z",
     "shell.execute_reply.started": "2025-10-18T00:07:22.323761Z"
    },
    "jupyter": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "def deepspeed_train_t5(num_samples: str, model_name: str, bucket: str):\n",
    "    import os\n",
    "    import time\n",
    "    import boto3\n",
    "    import torch\n",
    "    import torch.distributed as dist\n",
    "    from torch.utils.data.distributed import DistributedSampler\n",
    "    from transformers import T5Tokenizer, T5ForConditionalGeneration\n",
    "    from datasets import load_dataset\n",
    "    import deepspeed\n",
    "    import numpy as np\n",
    "\n",
    "    # Initialize distributed environment.\n",
    "    deepspeed.init_distributed(dist_backend=\"nccl\")\n",
    "    local_rank = int(os.environ[\"LOCAL_RANK\"])\n",
    "\n",
    "    # Define the summarization dataset class\n",
    "    class SummarizationDataset(torch.utils.data.Dataset):\n",
    "        def __init__(self, tokenizer, num_samples):\n",
    "            # Load publicly available CNN/DailyMail dataset\n",
    "            self.dataset = load_dataset(\"cnn_dailymail\", \"3.0.0\", split=f\"train[:{num_samples}]\")\n",
    "            self.tokenizer = tokenizer\n",
    "\n",
    "        def __len__(self):\n",
    "            return len(self.dataset)\n",
    "\n",
    "        def clean_text(self, text):\n",
    "            if text is None:\n",
    "                return \"\"\n",
    "            return text.replace(\"\\n\", \" \").replace(\"``\", \"\").replace('\"', \"\").strip()\n",
    "\n",
    "        def convert_to_features(self, example_batch):\n",
    "            # Map CNN/DailyMail fields to 'text' and 'summary'\n",
    "            input_ = self.clean_text(example_batch[\"article\"])\n",
    "            target_ = self.clean_text(example_batch[\"highlights\"])\n",
    "\n",
    "            source = self.tokenizer(\n",
    "                input_,\n",
    "                max_length=512,\n",
    "                padding=\"max_length\",\n",
    "                truncation=True,\n",
    "                return_tensors=\"pt\",\n",
    "            )\n",
    "            targets = self.tokenizer(\n",
    "                target_,\n",
    "                max_length=150,\n",
    "                padding=\"max_length\",\n",
    "                truncation=True,\n",
    "                return_tensors=\"pt\",\n",
    "            )\n",
    "\n",
    "            return source, targets\n",
    "\n",
    "        def __getitem__(self, index):\n",
    "            source, targets = self.convert_to_features(self.dataset[index])\n",
    "            return {\n",
    "                \"source_ids\": source[\"input_ids\"].squeeze(),\n",
    "                \"source_mask\": source[\"attention_mask\"].squeeze(),\n",
    "                \"target_ids\": targets[\"input_ids\"].squeeze(),\n",
    "                \"target_mask\": targets[\"attention_mask\"].squeeze(),\n",
    "            }\n",
    "\n",
    "    # Download model and tokenizer.\n",
    "    if dist.get_rank() == 0:\n",
    "        print(\"-\" * 100)\n",
    "        print(\"Downloading T5 Model\")\n",
    "        print(\"-\" * 100)\n",
    "\n",
    "    model = T5ForConditionalGeneration.from_pretrained(model_name)\n",
    "    tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
    "\n",
    "    # Download dataset.\n",
    "    dataset = SummarizationDataset(tokenizer, num_samples=int(num_samples))\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        dataset, batch_size=4, sampler=DistributedSampler(dataset)\n",
    "    )\n",
    "\n",
    "    # Define DeepSpeed configuration.\n",
    "    # Train batch size = micro batch size * gradient steps * GPUs (e.g. 2 x 1 x 8 = 16).\n",
    "    ds_config = {\n",
    "        \"train_micro_batch_size_per_gpu\": 2,\n",
    "        \"gradient_accumulation_steps\": 1,\n",
    "        # \"fp16\": {\"enabled\": True}, # If your GPU (e.g. V100) doesn't support bf16, use fp16.\n",
    "        \"bf16\": {\"enabled\": True}, # Enable mixed precision.\n",
    "        \"optimizer\": {\n",
    "            \"type\": \"AdamW\",\n",
    "            \"params\": {\"lr\": 0.002},\n",
    "        },\n",
    "        \"scheduler\": {\n",
    "            \"type\": \"WarmupLR\",\n",
    "            \"params\": {\n",
    "                \"warmup_min_lr\": 0,\n",
    "                \"warmup_max_lr\": 0.001,\n",
    "                \"warmup_num_steps\": 1000,\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "\n",
    "    # Initialize model with DeepSpeed.\n",
    "    model, _, _, _ = deepspeed.initialize(\n",
    "        config=ds_config,\n",
    "        model=model,\n",
    "        model_parameters=model.parameters(),\n",
    "    )\n",
    "\n",
    "    # Start training process.\n",
    "    if dist.get_rank() == 0:\n",
    "        print(\"-\" * 100)\n",
    "        print(\"Starting DeepSpeed distributed training...\")\n",
    "        print(\"-\" * 100)\n",
    "\n",
    "    t0 = time.time()\n",
    "    for epoch in range(1, 3):\n",
    "        losses = []\n",
    "        for batch_idx, batch in enumerate(train_loader):\n",
    "            for key in batch.keys():\n",
    "                batch[key] = batch[key].to(local_rank)\n",
    "            # Forward pass.\n",
    "            output = model(\n",
    "                input_ids=batch[\"source_ids\"],\n",
    "                attention_mask=batch[\"source_mask\"],\n",
    "                labels=batch[\"target_ids\"],\n",
    "            )\n",
    "            loss = output.loss\n",
    "\n",
    "            # Run backpropagation.\n",
    "            model.backward(loss)\n",
    "            # Weight updates.\n",
    "            model.step()\n",
    "            losses.append(loss.item())\n",
    "\n",
    "            if batch_idx % 10 == 0 and dist.get_rank() == 0:\n",
    "                print(\n",
    "                    f\"Train Epoch: {epoch} [{batch_idx * len(batch)}/{len(train_loader.dataset)} \"\n",
    "                    f\"({100.0 * batch_idx / len(train_loader):.0f}%)]\\tLoss: {loss.item():.6f}\"\n",
    "                )\n",
    "\n",
    "        if dist.get_rank() == 0:\n",
    "            print(\"-\" * 100)\n",
    "            print(f\"Average Train Loss: {np.mean(losses):.4f}\")\n",
    "            print(\"-\" * 100)\n",
    "\n",
    "    if dist.get_rank() == 0:\n",
    "        print(\"-\" * 100)\n",
    "        print(f\"DeepSpeed training time: {int(time.time() - t0)} seconds\")\n",
    "        print(\"-\" * 100)\n",
    "\n",
    "        print(\"Exporting HuggingFace model to S3\")\n",
    "        MODEL_PATH = os.path.join(\"/home/mpiuser\", model_name)\n",
    "        model.module.save_pretrained(MODEL_PATH)\n",
    "        tokenizer.save_pretrained(MODEL_PATH)\n",
    "\n",
    "        s3_bucket = boto3.resource(\"s3\").Bucket(bucket)\n",
    "        for file in os.listdir(MODEL_PATH):\n",
    "            print(f\"Uploading file {os.path.join(MODEL_PATH, file)}\")\n",
    "            s3_bucket.upload_file(\n",
    "                os.path.join(MODEL_PATH, file), os.path.join(model_name, file)\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "875abc44-69dd-41d8-bb2a-436f730a2dd0",
   "metadata": {},
   "source": [
    "## List Available Kubeflow Trainer Runtimes\n",
    "\n",
    "\n",
    "Get available Kubeflow Trainer Runtimes with the `list_runtimes()` API.\n",
    "\n",
    "You can inspect Runtime details, including the name, framework, and available default devices on the single node.\n",
    "\n",
    "- Runtimes with **CustomTrainer**: You must write the training script within the function.\n",
    "\n",
    "- Runtimes with **BuiltinTrainer**: You can configure settings (e.g., LoRA Config) for LLM fine-tuning Job.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "51d8bc1d-8d9b-48f7-866f-c6ad4ad4241b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:07:23.458715Z",
     "iopub.status.busy": "2025-10-18T00:07:23.458171Z",
     "iopub.status.idle": "2025-10-18T00:07:25.173611Z",
     "shell.execute_reply": "2025-10-18T00:07:25.172579Z",
     "shell.execute_reply.started": "2025-10-18T00:07:23.458676Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: deepspeed-distributed, Framework: deepspeed, Trainer Type: CustomTrainer\n",
      "\n",
      "Runtime devices: gpu x 1\n"
     ]
    }
   ],
   "source": [
    "from kubeflow.trainer import TrainerClient, CustomTrainer\n",
    "\n",
    "for r in TrainerClient().list_runtimes():\n",
    "    if r.name == \"deepspeed-distributed\":\n",
    "        print(f\"Name: {r.name}, Framework: {r.trainer.framework}, Trainer Type: {r.trainer.trainer_type.value}\\n\")\n",
    "        print(f\"Runtime devices: {r.trainer.device} x {r.trainer.device_count}\")\n",
    "        deepspeed_runtime = r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "206c389d-9dcb-474f-b0f5-cc098de2e183",
   "metadata": {},
   "source": [
    "## Create TrainJob for Distributed Training\n",
    "\n",
    "Use the `train()` API to scale the training code across 2 Nodes and 8 GPUs.\n",
    "\n",
    "Don't forget to update **the S3 bucket** name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5afe265b-9e07-4f5a-b4fb-5dfac8a61a86",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:07:47.348953Z",
     "iopub.status.busy": "2025-10-18T00:07:47.348616Z",
     "iopub.status.idle": "2025-10-18T00:07:47.352241Z",
     "shell.execute_reply": "2025-10-18T00:07:47.351567Z",
     "shell.execute_reply.started": "2025-10-18T00:07:47.348924Z"
    }
   },
   "outputs": [],
   "source": [
    "MODEL_NAME = \"t5-base\"\n",
    "# BUCKET_NAME = \"TODO: add your bucket here\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "76dab189-f184-4e48-be74-f32c0dea675b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:07:47.711362Z",
     "iopub.status.busy": "2025-10-18T00:07:47.710791Z",
     "iopub.status.idle": "2025-10-18T00:07:48.927014Z",
     "shell.execute_reply": "2025-10-18T00:07:48.926354Z",
     "shell.execute_reply.started": "2025-10-18T00:07:47.711322Z"
    }
   },
   "outputs": [],
   "source": [
    "args = {\n",
    "    \"num_samples\": \"2000\",\n",
    "    \"model_name\": MODEL_NAME,\n",
    "    \"bucket\": BUCKET_NAME,\n",
    "}\n",
    "\n",
    "job_id = TrainerClient().train(\n",
    "    trainer=CustomTrainer(\n",
    "        func=deepspeed_train_t5,\n",
    "        func_args=args,\n",
    "        packages_to_install=[\"boto3\"], # Custom packages to install at runtime.\n",
    "        num_nodes=2,\n",
    "        resources_per_node={\n",
    "            \"gpu\": 4\n",
    "        },\n",
    "    ),\n",
    "    runtime=deepspeed_runtime,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6e1811dd-4bf2-40cf-ad35-35a87271eb21",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:07:48.929294Z",
     "iopub.status.busy": "2025-10-18T00:07:48.929009Z",
     "iopub.status.idle": "2025-10-18T00:07:48.934923Z",
     "shell.execute_reply": "2025-10-18T00:07:48.934364Z",
     "shell.execute_reply.started": "2025-10-18T00:07:48.929263Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'l82624f2928b'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Train API generates a random TrainJob id.\n",
    "job_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0864121-895f-4e5a-87b8-7ea2d92e6630",
   "metadata": {},
   "source": [
    "## Check the TrainJob Info\n",
    "\n",
    "Use the `list_jobs()` and `get_job()` APIs to get information about created TrainJob and its steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8a945930-6cfe-4388-8ab8-462474f3f21f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:07:48.935708Z",
     "iopub.status.busy": "2025-10-18T00:07:48.935465Z",
     "iopub.status.idle": "2025-10-18T00:07:50.687949Z",
     "shell.execute_reply": "2025-10-18T00:07:50.686902Z",
     "shell.execute_reply.started": "2025-10-18T00:07:48.935689Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TrainJob: l82624f2928b, Status: Created, Created at: 2025-10-18 00:07:48+00:00\n"
     ]
    }
   ],
   "source": [
    "for job in TrainerClient().list_jobs():\n",
    "    print(f\"TrainJob: {job.name}, Status: {job.status}, Created at: {job.creation_timestamp}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "75e8c741-6d31-49a2-8667-d38e40d62430",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:08:28.619725Z",
     "iopub.status.busy": "2025-10-18T00:08:28.619183Z",
     "iopub.status.idle": "2025-10-18T00:08:30.388855Z",
     "shell.execute_reply": "2025-10-18T00:08:30.386957Z",
     "shell.execute_reply.started": "2025-10-18T00:08:28.619679Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step: node-0, Status: Running, Devices: gpu x 4\n",
      "\n",
      "Step: node-1, Status: Running, Devices: gpu x 4\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# We execute mpirun command on node-0, which functions as the MPI Launcher node.\n",
    "for c in TrainerClient().get_job(name=job_id).steps:\n",
    "    print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4fa3ea2-bdf5-40ca-a9ec-cc948949ea59",
   "metadata": {},
   "source": [
    "## Get the TrainJob Logs\n",
    "\n",
    "Use the `get_job_logs()` API to retrieve the TrainJob logs.\n",
    "\n",
    "Since we distribute the dataset accross 8 GPUs (2 nodes x 4 GPUs), each rank processes `round(2000 / 8) = 250` samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6e630fd3-f061-4fea-8024-7bffcefb257c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:10:53.039108Z",
     "iopub.status.busy": "2025-10-18T00:10:53.038643Z",
     "iopub.status.idle": "2025-10-18T00:10:54.952499Z",
     "shell.execute_reply": "2025-10-18T00:10:54.951093Z",
     "shell.execute_reply.started": "2025-10-18T00:10:53.039069Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: Permanently added '[l82624f2928b-node-0-0.l82624f2928b]:2222' (ECDSA) to the list of known hosts.\n",
      "[2025-10-18 00:08:23,111] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "df: /home/mpiuser/.triton/autotune: No such file or directory\n",
      "[2025-10-18 00:08:24,295] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2025-10-18 00:08:24,754] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:24,767] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:24,767] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:25,593] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2025-10-18 00:08:25,910] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:25,923] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:25,923] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:27,422] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:27,435] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:27,435] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:30,514] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2025-10-18 00:08:31,987] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:32,000] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:32,000] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:40,410] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "df: /home/mpiuser/.triton/autotune: No such file or directory\n",
      "[2025-10-18 00:08:41,794] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2025-10-18 00:08:42,072] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:42,085] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:42,085] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:43,484] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:43,497] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:43,498] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:49,499] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2025-10-18 00:08:50,045] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2025-10-18 00:08:50,997] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:51,010] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:51,010] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:51,680] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False\n",
      "[2025-10-18 00:08:51,692] [INFO] [comm.py:821:init_distributed] cdb=None\n",
      "[2025-10-18 00:08:51,693] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=0, local_rank=0, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=2, local_rank=2, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=1, local_rank=1, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=3, local_rank=3, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:852:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=4, local_rank=0, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=6, local_rank=2, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=5, local_rank=1, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=7, local_rank=3, world_size=8, master_addr=240.56.142.151, master_port=29500\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Downloading T5 Model\n",
      "----------------------------------------------------------------------------------------------------\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "Generating train split: 100%|██████████| 287113/287113 [00:03<00:00, 77899.43 examples/s]\n",
      "Generating train split: 100%|██████████| 287113/287113 [00:03<00:00, 79219.92 examples/s]s]\n",
      "Generating validation split: 100%|██████████| 13368/13368 [00:00<00:00, 74240.68 examples/s]\n",
      "Generating validation split: 100%|██████████| 13368/13368 [00:00<00:00, 80800.92 examples/s]\n",
      "Generating test split: 100%|██████████| 11490/11490 [00:00<00:00, 81368.13 examples/s]\n",
      "[2025-10-18 00:09:10,600] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "[2025-10-18 00:09:10,609] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed info: version=0.17.4, git-hash=unknown, git-branch=unknown\n",
      "[2025-10-18 00:09:10,609] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "[2025-10-18 00:09:10,610] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "[2025-10-18 00:09:10,613] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "Generating test split: 100%|██████████| 11490/11490 [00:00<00:00, 82023.46 examples/s]\n",
      "[2025-10-18 00:09:10,728] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "[2025-10-18 00:09:10,738] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "[2025-10-18 00:09:10,746] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "[2025-10-18 00:09:10,758] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8\n",
      "[2025-10-18 00:09:11,696] [INFO] [engine.py:1339:_configure_distributed_model] ********** distributed groups summary **********\n",
      "\t self.dp_world_size=8\n",
      "\t self.mp_world_size=1\n",
      "\t self.seq_dp_world_size=8\n",
      "\t self.sequence_parallel_size=1\n",
      "***********************************************\n",
      "[2025-10-18 00:09:11,843] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Creating extension directory /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam...\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Creating extension directory /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam...\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...\n",
      "Detected CUDA files, patching ldflags\n",
      "Emitting ninja build file /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam/build.ninja...\n",
      "Detected CUDA files, patching ldflags\n",
      "Emitting ninja build file /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam/build.ninja...\n",
      "/home/mpiuser/.local/lib/python3.10/site-packages/torch/utils/cpp_extension.py:2356: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. \n",
      "If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].\n",
      "  warnings.warn(\n",
      "/home/mpiuser/.local/lib/python3.10/site-packages/torch/utils/cpp_extension.py:2356: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. \n",
      "If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].\n",
      "  warnings.warn(\n",
      "Building extension module fused_adam...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n",
      "Building extension module fused_adam...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n",
      "[1/3] /usr/local/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output multi_tensor_adam.cuda.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1016\\\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -O3 -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -lineinfo --use_fast_math -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_80,code=compute_80 -UC10_USE_GLOG -DBF16_AVAILABLE -U__CUDA_NO_BFLOAT16_OPERATORS__ -U__CUDA_NO_BFLOAT162_OPERATORS__ -U__CUDA_NO_BFLOAT16_CONVERSIONS__ -std=c++17 -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu -o multi_tensor_adam.cuda.o \n",
      "[1/3] /usr/local/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output multi_tensor_adam.cuda.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1016\\\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -O3 -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -lineinfo --use_fast_math -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_80,code=compute_80 -UC10_USE_GLOG -DBF16_AVAILABLE -U__CUDA_NO_BFLOAT16_OPERATORS__ -U__CUDA_NO_BFLOAT162_OPERATORS__ -U__CUDA_NO_BFLOAT16_CONVERSIONS__ -std=c++17 -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu -o multi_tensor_adam.cuda.o \n",
      "[2/3] c++ -MMD -MF fused_adam_frontend.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1016\\\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -O3 -std=c++17 -g -Wno-reorder -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -UC10_USE_GLOG -DBF16_AVAILABLE -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp -o fused_adam_frontend.o \n",
      "[2/3] c++ -MMD -MF fused_adam_frontend.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1016\\\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -O3 -std=c++17 -g -Wno-reorder -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -UC10_USE_GLOG -DBF16_AVAILABLE -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp -o fused_adam_frontend.o \n",
      "[3/3] c++ fused_adam_frontend.o multi_tensor_adam.cuda.o -shared -L/home/mpiuser/.local/lib/python3.10/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/cuda/lib64 -lcudart -o fused_adam.so\n",
      "Loading extension module fused_adam...\n",
      "Time to load fused_adam op: 25.939489603042603 seconds\n",
      "[2025-10-18 00:09:37,813] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "Loading extension module fused_adam...\n",
      "Time to load fused_adam op: 26.032996654510498 seconds\n",
      "Loading extension module fused_adam...\n",
      "Time to load fused_adam op: 26.032512187957764 seconds\n",
      "Loading extension module fused_adam...\n",
      "Time to load fused_adam op: 26.037198543548584 seconds\n",
      "[2025-10-18 00:09:37,883] [INFO] [logging.py:107:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adamw as basic optimizer\n",
      "[2025-10-18 00:09:37,883] [INFO] [logging.py:107:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer\n",
      "[2025-10-18 00:09:37,894] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed Basic Optimizer = FusedAdam\n",
      "[2025-10-18 00:09:37,894] [INFO] [logging.py:107:log_dist] [Rank 0] Creating BF16 optimizer\n",
      "[3/3] c++ fused_adam_frontend.o multi_tensor_adam.cuda.o -shared -L/home/mpiuser/.local/lib/python3.10/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/cuda/lib64 -lcudart -o fused_adam.so\n",
      "Loading extension module fused_adam...\n",
      "Time to load fused_adam op: 26.056629180908203 seconds\n",
      "[2025-10-18 00:09:37,908] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "[2025-10-18 00:09:37,915] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "[2025-10-18 00:09:37,954] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "Loading extension module fused_adam...\n",
      "Loading extension module fused_adam...\n",
      "Time to load fused_adam op: 26.13313388824463 seconds\n",
      "Time to load fused_adam op: 26.13256287574768 seconds\n",
      "Loading extension module fused_adam...\n",
      "Time to load fused_adam op: 26.137150526046753 seconds\n",
      "[2025-10-18 00:09:38,010] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "[2025-10-18 00:09:38,011] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "[2025-10-18 00:09:38,018] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "[2025-10-18 00:09:38,120] [INFO] [utils.py:781:see_memory_usage] begin bf16_optimizer\n",
      "[2025-10-18 00:09:38,121] [INFO] [utils.py:782:see_memory_usage] MA 0.43 GB         Max_MA 0.43 GB         CA 0.44 GB         Max_CA 0 GB \n",
      "[2025-10-18 00:09:38,121] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 23.24 GB, percent = 2.1%\n",
      "[2025-10-18 00:09:38,314] [INFO] [utils.py:781:see_memory_usage] before initializing group 0\n",
      "[2025-10-18 00:09:38,315] [INFO] [utils.py:782:see_memory_usage] MA 0.43 GB         Max_MA 0.43 GB         CA 0.44 GB         Max_CA 0 GB \n",
      "[2025-10-18 00:09:38,315] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 24.34 GB, percent = 2.2%\n",
      "[2025-10-18 00:09:38,523] [INFO] [utils.py:781:see_memory_usage] after initializing group 0\n",
      "[2025-10-18 00:09:38,524] [INFO] [utils.py:782:see_memory_usage] MA 1.35 GB         Max_MA 1.35 GB         CA 1.85 GB         Max_CA 2 GB \n",
      "[2025-10-18 00:09:38,524] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 25.03 GB, percent = 2.2%\n",
      "[2025-10-18 00:09:38,704] [INFO] [utils.py:781:see_memory_usage] end bf16_ optimizer\n",
      "[2025-10-18 00:09:38,705] [INFO] [utils.py:782:see_memory_usage] MA 1.35 GB         Max_MA 1.35 GB         CA 1.85 GB         Max_CA 2 GB \n",
      "[2025-10-18 00:09:38,705] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 25.14 GB, percent = 2.2%\n",
      "[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed Final Optimizer = BF16_Optimizer\n",
      "[2025-10-18 00:09:38,705] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started\n",
      "[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed using configured LR scheduler = WarmupLR\n",
      "[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed LR Scheduler = <deepspeed.runtime.lr_schedules.WarmupLR object at 0x7f91ac245fc0>\n",
      "[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] step=0, skipped=0, lr=[0], mom=[(0.9, 0.999)]\n",
      "[2025-10-18 00:09:38,706] [INFO] [logging.py:107:log_dist] [Rank 0] [TorchCheckpointEngine] Initialized with serialization = True\n",
      "[2025-10-18 00:09:38,706] [INFO] [config.py:954:print] DeepSpeedEngine configuration:\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   activation_checkpointing_config  {\n",
      "    \"partition_activations\": false, \n",
      "    \"contiguous_memory_optimization\": false, \n",
      "    \"cpu_checkpointing\": false, \n",
      "    \"number_checkpoints\": null, \n",
      "    \"synchronize_checkpoint_boundary\": false, \n",
      "    \"profile\": false\n",
      "}\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'intra_op_parallelism': 1, 'single_submit': False, 'overlap_events': True, 'use_gds': False}\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   amp_enabled .................. False\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   amp_params ................... False\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   autotuning_config ............ {\n",
      "    \"enabled\": false, \n",
      "    \"start_step\": null, \n",
      "    \"end_step\": null, \n",
      "    \"metric_path\": null, \n",
      "    \"arg_mappings\": null, \n",
      "    \"metric\": \"throughput\", \n",
      "    \"model_info\": null, \n",
      "    \"results_dir\": \"autotuning_results\", \n",
      "    \"exps_dir\": \"autotuning_exps\", \n",
      "    \"overwrite\": true, \n",
      "    \"fast\": true, \n",
      "    \"start_profile_step\": 3, \n",
      "    \"end_profile_step\": 5, \n",
      "    \"tuner_type\": \"gridsearch\", \n",
      "    \"tuner_early_stopping\": 5, \n",
      "    \"tuner_num_trials\": 50, \n",
      "    \"model_info_path\": null, \n",
      "    \"mp_size\": 1, \n",
      "    \"max_train_batch_size\": null, \n",
      "    \"min_train_batch_size\": 1, \n",
      "    \"max_train_micro_batch_size_per_gpu\": 1.024000e+03, \n",
      "    \"min_train_micro_batch_size_per_gpu\": 1, \n",
      "    \"num_tuning_micro_batch_sizes\": 3\n",
      "}\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   bfloat16_config .............. enabled=True immediate_grad_update=False check_grad_overflow=False\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_config ............ {'tag_validation': 'WARN', 'checkpoint_serialization': True, 'writer': None}\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_parallel_write_pipeline  False\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_tag_validation_enabled  True\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_tag_validation_fail  False\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7f91ac289630>\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   communication_data_type ...... None\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   compile_config ............... deepcompile=False free_activation=False offload_activation=False offload_opt_states=False double_buffer=True symmetric_memory=False debug_log=False offload_parameters=False sync_before_reduce=False sync_after_reduce=False sync_before_allgather=False sync_after_allgather=False keep_int_input_tensors=True keep_all_input_tensors=False\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}\n",
      "[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   curriculum_enabled_legacy .... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   curriculum_params_legacy ..... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'pin_memory': False, 'curriculum_learning': {'enabled': False}, 'dynamic_batching': {'enabled': False, 'lr_scaling_method': 'linear', 'min_batch_size': 1, 'max_batch_size': None, 'sequence_picking_order': 'dataloader', 'verbose': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   data_efficiency_enabled ...... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   dataloader_drop_last ......... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   disable_allgather ............ False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   dump_state ................... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_enabled ........... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_gas_boundary_resolution  1\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_layer_name ........ bert.encoder.layer\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_layer_num ......... 0\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_max_iter .......... 100\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_stability ......... 1e-06\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_tol ............... 0.01\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_verbose ........... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   elasticity_enabled ........... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   float16_config ............... enabled=False auto_cast=False loss_scale=0.0 initial_scale_power=16 loss_scale_window=1000 hysteresis=2 consecutive_hysteresis=False min_loss_scale=1 fp16_master_weights_and_grads=False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   flops_profiler_config ........ {\n",
      "    \"enabled\": false, \n",
      "    \"recompute_fwd_factor\": 0.0, \n",
      "    \"profile_step\": 1, \n",
      "    \"module_depth\": -1, \n",
      "    \"top_modules\": 1, \n",
      "    \"detailed\": true, \n",
      "    \"output_file\": null\n",
      "}\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   global_rank .................. 0\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   grad_accum_dtype ............. None\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   gradient_accumulation_steps .. 1\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   gradient_clipping ............ 0.0\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   gradient_predivide_factor .... 1.0\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   graph_harvesting ............. False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   load_universal_checkpoint .... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   memory_breakdown ............. False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   mics_hierarchial_params_gather  False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   mics_shard_size .............. -1\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') comet=CometConfig(enabled=False, samples_log_interval=100, project=None, workspace=None, api_key=None, experiment_name=None, experiment_key=None, online=None, mode=None) wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName')\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   nebula_config ................ {\n",
      "    \"enabled\": false, \n",
      "    \"persistent_storage_path\": null, \n",
      "    \"persistent_time_interval\": 100, \n",
      "    \"num_of_version_in_retention\": 2, \n",
      "    \"enable_nebula_load\": true, \n",
      "    \"load_path\": null\n",
      "}\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   optimizer_legacy_fusion ...... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   optimizer_name ............... adamw\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   optimizer_params ............. {'lr': 0.002}\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0, 'pipe_partitioned': True, 'grad_partitioned': True}\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   pld_enabled .................. False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   pld_params ................... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   prescale_gradients ........... False\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   scheduler_name ............... WarmupLR\n",
      "[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   scheduler_params ............. {'warmup_min_lr': 0, 'warmup_max_lr': 0.001, 'warmup_num_steps': 1000}\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   seq_parallel_communication_data_type  torch.float32\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   sparse_attention ............. None\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   sparse_gradients_enabled ..... False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   steps_per_print .............. None\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   tensor_parallel_config ....... dtype=torch.float16 autotp_size=0 tp_overlap_comm=False tensor_parallel=TPConfig(tp_size=1, tp_grain_size=1, mpu=None, tp_group=None) injection_policy_tuple=None keep_module_on_host=False replace_with_kernel_inject=False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   timers_config ................ enabled=True synchronized=True\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   torch_autocast_dtype ......... None\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   torch_autocast_enabled ....... False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   torch_autocast_lower_precision_safe_modules  None\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   train_batch_size ............. 16\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   train_micro_batch_size_per_gpu  2\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   use_data_before_expert_parallel_  False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   use_node_local_storage ....... False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   wall_clock_breakdown ......... False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   weight_quantization_config ... None\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   world_size ................... 8\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_allow_untested_optimizer  False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_config .................. stage=0 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=500000000 use_multi_rank_bucket_allreduce=True allgather_partitions=True allgather_bucket_size=500000000 overlap_comm=False load_from_fp32_weights=True elastic_checkpoint=False offload_param=None offload_optimizer=None sub_group_size=1000000000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=50000000 param_persistence_threshold=100000 model_persistence_threshold=9223372036854775807 max_live_parameters=1000000000 max_reuse_distance=1000000000 gather_16bit_weights_on_model_save=False module_granularity_threshold=0 use_all_reduce_for_fetch_params=False stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_nontrainable_weights=False zero_quantized_gradients=False zeropp_loco_param=None mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True pipeline_loading_checkpoint=False override_module_apply=True log_trace_cache_warnings=False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_enabled ................. False\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_force_ds_cpu_optimizer .. True\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_optimization_stage ...... 0\n",
      "[2025-10-18 00:09:38,709] [INFO] [config.py:944:print_user_config]   json = {\n",
      "    \"train_micro_batch_size_per_gpu\": 2, \n",
      "    \"gradient_accumulation_steps\": 1, \n",
      "    \"bf16\": {\n",
      "        \"enabled\": true\n",
      "    }, \n",
      "    \"optimizer\": {\n",
      "        \"type\": \"AdamW\", \n",
      "        \"params\": {\n",
      "            \"lr\": 0.002\n",
      "        }\n",
      "    }, \n",
      "    \"scheduler\": {\n",
      "        \"type\": \"WarmupLR\", \n",
      "        \"params\": {\n",
      "            \"warmup_min_lr\": 0, \n",
      "            \"warmup_max_lr\": 0.001, \n",
      "            \"warmup_num_steps\": 1000\n",
      "        }\n",
      "    }\n",
      "}\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Starting DeepSpeed distributed training...\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Train Epoch: 1 [0/2000 (0%)]\tLoss: 10.250000\n",
      "Train Epoch: 1 [40/2000 (16%)]\tLoss: 0.878906\n",
      "Train Epoch: 1 [80/2000 (32%)]\tLoss: 0.960938\n",
      "Train Epoch: 1 [120/2000 (48%)]\tLoss: 0.957031\n",
      "Train Epoch: 1 [160/2000 (63%)]\tLoss: 0.707031\n",
      "Train Epoch: 1 [200/2000 (79%)]\tLoss: 0.621094\n",
      "Train Epoch: 1 [240/2000 (95%)]\tLoss: 0.761719\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Average Train Loss: 1.5145\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Train Epoch: 2 [0/2000 (0%)]\tLoss: 0.625000\n",
      "Train Epoch: 2 [40/2000 (16%)]\tLoss: 0.558594\n",
      "Train Epoch: 2 [80/2000 (32%)]\tLoss: 0.617188\n",
      "Train Epoch: 2 [120/2000 (48%)]\tLoss: 0.667969\n",
      "Train Epoch: 2 [160/2000 (63%)]\tLoss: 0.494141\n",
      "Train Epoch: 2 [200/2000 (79%)]\tLoss: 0.365234\n",
      "Train Epoch: 2 [240/2000 (95%)]\tLoss: 0.523438\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Average Train Loss: 0.5215\n",
      "----------------------------------------------------------------------------------------------------\n",
      "----------------------------------------------------------------------------------------------------\n",
      "DeepSpeed training time: 59 seconds\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Exporting HuggingFace model to S3\n",
      "Uploading file /home/mpiuser/t5-base/config.json\n",
      "Uploading file /home/mpiuser/t5-base/generation_config.json\n",
      "Uploading file /home/mpiuser/t5-base/model.safetensors\n",
      "[rank4]:[W1018 00:10:39.692111139 ProcessGroupNCCL.cpp:1479] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())\n",
      "Uploading file /home/mpiuser/t5-base/tokenizer_config.json\n",
      "Uploading file /home/mpiuser/t5-base/special_tokens_map.json\n",
      "Uploading file /home/mpiuser/t5-base/added_tokens.json\n",
      "Uploading file /home/mpiuser/t5-base/spiece.model\n",
      "[rank0]:[W1018 00:10:44.046274646 ProcessGroupNCCL.cpp:1479] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())\n"
     ]
    }
   ],
   "source": [
    "for logline in TrainerClient().get_job_logs(job_id, follow=True):\n",
    "    print(logline)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "030ea83c-9b1c-477a-a3f3-72d659066678",
   "metadata": {},
   "source": [
    "## Download the Trained Model\n",
    "\n",
    "Finally, download fine-tuned model from S3 for evaluations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "def0bfea-9fea-4616-bf43-d945d768e630",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:13:47.447025Z",
     "iopub.status.busy": "2025-10-18T00:13:47.446531Z",
     "iopub.status.idle": "2025-10-18T00:13:47.451244Z",
     "shell.execute_reply": "2025-10-18T00:13:47.450121Z",
     "shell.execute_reply.started": "2025-10-18T00:13:47.446992Z"
    }
   },
   "outputs": [],
   "source": [
    "LOCAL_DIR = \"./t5-base\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b11d5623-3425-41ea-87e7-ec6992540994",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:13:48.301747Z",
     "iopub.status.busy": "2025-10-18T00:13:48.301256Z",
     "iopub.status.idle": "2025-10-18T00:14:15.215668Z",
     "shell.execute_reply": "2025-10-18T00:14:15.215088Z",
     "shell.execute_reply.started": "2025-10-18T00:13:48.301714Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading file: t5-base/added_tokens.json\n",
      "Downloading file: t5-base/config.json\n",
      "Downloading file: t5-base/generation_config.json\n",
      "Downloading file: t5-base/model.safetensors\n",
      "Downloading file: t5-base/special_tokens_map.json\n",
      "Downloading file: t5-base/spiece.model\n",
      "Downloading file: t5-base/tokenizer_config.json\n"
     ]
    }
   ],
   "source": [
    "import boto3\n",
    "import os\n",
    "\n",
    "os.makedirs(LOCAL_DIR, exist_ok=True)\n",
    "s3 = boto3.client(\"s3\")\n",
    "for obj in s3.list_objects_v2(Bucket=BUCKET_NAME, Prefix=\"t5-base\")[\"Contents\"]:\n",
    "    file = obj[\"Key\"]\n",
    "\n",
    "    print(f\"Downloading file: {file}\")\n",
    "    s3.download_file(BUCKET_NAME, file, os.path.join(LOCAL_DIR, os.path.basename(file)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3075e5d9-62ec-495e-9bc8-18b5b3269e25",
   "metadata": {},
   "source": [
    "## Evaluate Fine-Tuned T5 Model\n",
    "\n",
    "After model is downloaded, you can load it into the HuggingFace pipeline.\n",
    "\n",
    "The T5 model performs well for NLP tasks such as summarization, translation, and text classification.\n",
    "\n",
    "In the example below, we'll demonstrate how to use a fine-tuned version of the T5 model to summarize documentation related to the Kubeflow Trainer project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e54caa0d-bd5a-4c3d-862b-3041072461b7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-10-18T00:14:42.497494Z",
     "iopub.status.busy": "2025-10-18T00:14:42.497009Z",
     "iopub.status.idle": "2025-10-18T00:14:43.847692Z",
     "shell.execute_reply": "2025-10-18T00:14:43.847466Z",
     "shell.execute_reply.started": "2025-10-18T00:14:42.497459Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Device set to use mps:0\n",
      "Both `max_new_tokens` (=256) and `max_length`(=100) seem to have been set. `max_new_tokens` will take precedence. Please refer to the documentation for more information. (https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'summary_text': 'Kubeflow Trainer allows you to fine-tune and enabling scalable, distributed training of machine learning (ML) models . Kubernetes-native project designed for large language models (LLMs).'}]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, pipeline\n",
    "\n",
    "# Load the fine-tuned T5 model.\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_DIR)\n",
    "tokenizer = AutoTokenizer.from_pretrained(LOCAL_DIR)\n",
    "\n",
    "summarizer = pipeline(\"summarization\", model=model, tokenizer=tokenizer, framework=\"pt\")\n",
    "\n",
    "text = \"\"\"\n",
    "summarize: In Kubeflow Trainer you can integrate other ML libraries such as HuggingFace,\n",
    "DeepSpeed, or Megatron-LM with Kubeflow Trainer to orchestrate their ML training on Kubernetes.\n",
    "Kubeflow Trainer allows you to effortlessly develop your LLMs with the Kubeflow Python SDK\n",
    "and build Kubernetes-native Training Runtimes with Kubernetes Custom Resources APIs.\n",
    "Kubeflow Trainer is a Kubernetes-native project designed for large language models (LLMs)\n",
    "fine-tuning and enabling scalable, distributed training of machine learning (ML)\n",
    "models across various frameworks, including PyTorch, JAX, TensorFlow, and XGBoost.\n",
    "\"\"\"\n",
    "\n",
    "summarizer(text, min_length=5, max_length=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54c68c68-9613-400c-9937-23b54d4d4b1d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
